import numpy as np
from PyQt5 import QtCore, QtGui, QtWidgets
from ezgl.items import *
from ezgl import GLViewWidget, tb, Matrix4x4, Quaternion
from pathlib import Path
import pickle

BASE_DIR = Path(__file__).resolve().parent
np.set_printoptions(precision=3, suppress=True)

class GLView(GLViewWidget):
    def __init__(self, parent=None, **kwargs):
        super().__init__(parent=parent, **kwargs)
        self.setup_ui()

        with open ('obj_2.urdf.pkl', 'rb') as f:
            data = pickle.load(f)
        print(data.keys())

        self.obj_pose = []
        obj_pose = data['obj_pose']
        for objp in obj_pose:
            self.obj_pose.extend(objp)

        self.franka_dof_pos = []
        franka_dof_pos = data['franka_dof_pos']
        for franka_dof in franka_dof_pos:
            self.franka_dof_pos.extend(franka_dof)

        self.base_frame = data['base_frame'].cpu().numpy()
        self.base_frame = Matrix4x4.fromVector7d(self.base_frame[0], self.base_frame[1], self.base_frame[2], self.base_frame[6],
                                                 self.base_frame[3], self.base_frame[4], self.base_frame[5])
        self.cnt = 0

    def setup_ui(self):
        self.resize(1200, 900)
        self.camera.set_vector6d((0.0, 0.0, -3.8, 0, -70, -90))
        self.ax = GLAxisItem(fix_to_corner=True)

        # -- lights
        light1 = PointLight(pos=(-2, -2, 3), ambient=(0., 0., 0.), diffuse=(0.6, 0.6, 0.6), visible=True, directional=True, render_shadow=False)
        light2 = PointLight(pos=(5, 3, 4), ambient=(0.2, 0.2, 0.2), diffuse=(0.5, 0.5, 0.5), visible=True, directional=True, render_shadow=False)
        self.light = PointLight(pos=(5, 7, 10), ambient=(0.6, 0.6, 0.6), diffuse=(0.6, 0.6, 0.6), visible=False, directional=True, render_shadow=False)
        self.light.setShadow(-3, 3, -3, 3, -3, 3, bias=0.05)

        # -- grid
        self.grid = GLGridItem(
            size=(11, 11), spacing=(0.5, 0.5), lineWidth=1, color=np.array([0.78, 0.71, 0.60])*1.5,
            lineColor=(0.4, 0.3, 0.2), lights=[self.light]
        ).rotate(90, 1, 0, 0)

        self.model = GLURDFItem("./assets/panda/panda_with_finger.urdf", lights=[light1, light2], glOptions="translucent_cull")
        self.obj = GLModelItem(path="./assets/obj/DeutschConnector.STL", lights=[light1, light2], glOptions="translucent_cull")
        self.obj.addChildItem(GLAxisItem(size=(0.12, 0.12, 0.12), tip_size=0.12))

        self.addItems([self.ax, self.grid, self.obj, self.model])

        # tool_box
        links_name = self.model.get_links_name()

        with tb.window("control", self, 10, size=(400, 300), pos=(-400, 0)):
            tb.add_timer("timer", interval_ms=30, callback=self.onTimeout)
            with tb.group("axis_visible", horizontal=False, collapsed=True, collapsible=True):
                tb.add_checklist("axis_visible", items=links_name, callback=self.on_changed,
                                exclusive=False, horizontal=False, value=None)
            tb.add_button("next", value=False, callback=self.next_pose)

    def next_pose(self):
        self.cnt = (self.cnt + 1) % len(self.obj_pose)
        # 设置物体位姿
        pose = self.obj_pose[self.cnt]
        pose = Matrix4x4.fromVector7d(pose[0], pose[1], pose[2], pose[6], pose[3], pose[4], pose[5])
        self.obj.setTransform(self.base_frame.inverse() * pose)

        # 设置机器人
        joints = self.franka_dof_pos[self.cnt]
        self.model.set_joints(joints)

    def onTimeout(self):
        self.update()

    def on_changed(self, data):
        tag, val = data
        label = self.sender().get_label()
        if label == "axis_visible":
            self.model.set_link(tag, axis_visiable=val)


if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)
    win = GLView(None)
    win.show()
    sys.exit(app.exec_())